#!/bin/env python
# -*- coding: utf-8 -*-
"""
演示使用内核缓存管理器优化Taichi编译时间

该示例展示如何使用内核缓存管理器来避免重复编译Taichi内核，
从而提高多次运行程序的性能。
"""

import os
import time
import taichi as ti
import sys
import numpy as np
from taichi_dither.kernel_cache_manager import enable_kernel_cache
from taichi_dither import DitherFactory, DitherAlgorithm, BayerDitherer

def run_dither_test(ditherer, img_path, num_colors, use_color=True, width=512):
    """
    运行抖动算法测试
    
    :param ditherer: 抖动算法实例
    :param img_path: 输入图像路径
    :param num_colors: 颜色数量
    :param use_color: 是否使用彩色
    :param width: 输出宽度
    :return: 处理后的图像
    """
    start_time = time.time()
    
    if isinstance(ditherer, BayerDitherer):
        processed_img = ditherer.process_image(
            img_path, 
            num_colors=num_colors,
            bayer_n=4,  # Bayer矩阵阶数
            use_color=use_color,
            output_width=width
        )
    else:
        processed_img = ditherer.process_image(
            img_path, 
            num_colors=num_colors,
            use_color=use_color,
            output_width=width
        )
    
    end_time = time.time()
    return processed_img, end_time - start_time

def main():
    """主函数：演示内核缓存的效果"""
    # 获取脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    img_path = os.path.join(script_dir, 'test.jpg')
    
    # 确保测试图像存在
    if not os.path.exists(img_path):
        print(f"错误：测试图像不存在: {img_path}")
        print("请在examples目录中放置一个名为test.jpg的测试图像")
        return
    
    # 设置缓存参数
    print("初始化Taichi环境，启用内核缓存...")
    cache_manager = enable_kernel_cache(
        arch=ti.gpu,  # 使用GPU后端
        debug=False,
        kernel_profiler=True,
        offline_cache=True  # 启用离线缓存
    )
    
    # 测试配置
    test_configs = [
        {"name": DitherAlgorithm.FLOYD, "num_colors": 8, "use_color": True, "width": 512},
        {"name": DitherAlgorithm.BAYER, "num_colors": 8, "use_color": True, "width": 512},
        {"name": DitherAlgorithm.ATKINSON, "num_colors": 8, "use_color": True, "width": 512}
    ]
    
    # 使用工厂类创建抖动算法实例
    algorithms = {}
    for config in test_configs:
        algo = config["name"]
        algorithms[algo] = DitherFactory.create_ditherer(
            algorithm=algo,
            arch=ti.gpu, 
            debug=False
        )
    
    print("\n===== 第一次运行（包含编译时间）=====")
    first_run_times = {}
    
    # 首次运行（包括编译时间）
    for config in test_configs:
        algo = config["name"]
        algo_name = algo.value if isinstance(algo, DitherAlgorithm) else algo
        print(f"\n测试算法: {algo_name}")
        
        ditherer = algorithms[algo]
        processed_img, runtime = run_dither_test(
            ditherer,
            img_path,
            config["num_colors"],
            config["use_color"],
            config["width"]
        )
        
        first_run_times[algo] = runtime
        print(f"处理耗时: {runtime:.4f} 秒")
        
        # 保存处理后的图像
        # output_filename = f"output_cached_{algo_name}_{config['num_colors']}colors.jpg"
        # ditherer.save_image(processed_img, output_filename)
        # print(f"已保存到: {output_filename}")
    
    print("\n===== 第二次运行（应使用缓存） =====")
    second_run_times = {}
    
    # 第二次运行（应使用缓存）
    for config in test_configs:
        algo = config["name"]
        algo_name = algo.value if isinstance(algo, DitherAlgorithm) else algo
        print(f"\n测试算法: {algo_name}")
        
        ditherer = algorithms[algo]
        processed_img, runtime = run_dither_test(
            ditherer,
            img_path,
            config["num_colors"],
            config["use_color"],
            config["width"]
        )
        
        second_run_times[algo] = runtime
        print(f"处理耗时: {runtime:.4f} 秒")
        improvement = (first_run_times[algo] - runtime) / first_run_times[algo] * 100
        print(f"性能提升: {improvement:.2f}%")
    
    # 打印汇总信息
    print("\n===== 性能对比汇总 =====")
    print("算法名称   | 首次运行(秒) | 二次运行(秒) | 提升比例")
    print("----------|------------|------------|--------")
    
    for algo in first_run_times.keys():
        algo_name = algo.value if isinstance(algo, DitherAlgorithm) else algo
        time1 = first_run_times[algo]
        time2 = second_run_times[algo]
        improvement = (time1 - time2) / time1 * 100
        print(f"{algo_name:<10} | {time1:<12.4f} | {time2:<12.4f} | {improvement:<.2f}%")
    
    # 显示缓存统计
    stats = cache_manager.cache_stats()
    print("\n===== 缓存统计 =====")
    print(f"缓存文件数: {stats['file_count']}")
    print(f"缓存总大小: {stats['total_size_mb']:.2f} MB")
    print(f"命中次数: {stats['hit_count']}")
    print(f"未命中次数: {stats['miss_count']}")
    print(f"命中率: {stats['hit_rate']:.2f}")
    
    # 打印Taichi内核性能信息
    if hasattr(ti, "profiler") and hasattr(ti.profiler, "print_kernel_profiler_info"):
        print("\n===== Taichi内核性能分析 =====")
        ti.profiler.print_kernel_profiler_info()
    elif hasattr(ti, "summary"):
        print("\n===== Taichi内核性能分析 =====")
        ti.summary()
    
    print("\n程序执行完成。")

if __name__ == "__main__":
    main() 